package database

import (
	"context"
	"github.com/go-xorm/xorm"
)

const transactionKey = "TRANS"

type SupportedTransaction interface {
	NewTransaction() ITransaction
	GetConn() xorm.Interface
}

func TryGetSession(ctx context.Context, repo SupportedTransaction) xorm.Interface {
	var session xorm.Interface
	if IsTransaction(ctx) {
		session = GetTransaction(ctx, repo).(*xorm.Session)
	} else {
		session = repo.GetConn()
	}
	return session
}

type ITransaction interface {
	Rollback() error
	Commit() error
}

type CtxTransaction struct {
	Tx ITransaction
}

func (ctx *CtxTransaction) Rollback() error {
	return ctx.Tx.Rollback()
}

func (ctx *CtxTransaction) Commit() error {
	return ctx.Tx.Commit()
}

func WithTransaction(ctx context.Context) context.Context {
	return context.WithValue(ctx, transactionKey, &CtxTransaction{})
}

func IsTransaction(ctx context.Context) bool {
	return ctx.Value(transactionKey) != nil
}

func GetTransaction(ctx context.Context, conn SupportedTransaction) ITransaction {
	tx := ctx.Value(transactionKey).(*CtxTransaction)
	if tx.Tx == nil && conn != nil {
		tx.Tx = conn.NewTransaction()
	}
	return ctx.Value(transactionKey).(*CtxTransaction).Tx
}

func Rollback(ctx context.Context) error {
	tx := GetTransaction(ctx, nil)
	if tx != nil {
		return tx.(ITransaction).Rollback()
	}
	return nil
}

func Commit(ctx context.Context) error {
	tx := GetTransaction(ctx, nil)
	if tx != nil {
		return tx.(ITransaction).Commit()
	}
	return nil
}
